// Fill out your copyright notice in the Description page of Project Settings.


#include "ScreenShot/Subsystems/ScreenShotSubsystem.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "qrencode.h"
#include "Engine/Texture2D.h"
#include "RenderUtils.h"
#include "Internationalization/Regex.h"

FDelegateHandle UScreenShotSubsystem::DelegateHandle = {};

void UScreenShotSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);
}

void UScreenShotSubsystem::Deinitialize()
{
	Super::Deinitialize();
}

bool UScreenShotSubsystem::ShouldCreateSubsystem(UObject* Outer) const
{
	TArray<UClass*> ChildClasses;
	GetDerivedClasses(GetClass(), ChildClasses, false);
	return ChildClasses.Num() == 0;
}

UTexture2D* UScreenShotSubsystem::CreateTexture2DFromBytes(UPARAM(ref)TArray<uint8>& ImageData) const
{
		EImageFormat ImageFormat = EImageFormat::Invalid;
			if (ImageData.Num() >= 8 && ImageData[0] == 0x89 && ImageData[1] == 0x50 && ImageData[2] == 0x4E && ImageData[3] == 0x47)
			{
				ImageFormat = EImageFormat::PNG;
			}
			else if (ImageData[0] == 0xFF && ImageData[1] == 0xD8 && ImageData[2] == 0xFF)
			{
				ImageFormat = EImageFormat::JPEG;
			}
			IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(
				FName("ImageWrapper"));
			TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(ImageFormat);
			if (ImageWrapper->SetCompressed(ImageData.GetData(), ImageData.Num()))
			{
				TArray<uint8> UncompressedRGBA;
				if (ImageWrapper->GetRaw(ERGBFormat::RGBA, 8, UncompressedRGBA))
				{
					TObjectPtr<UTexture2D> Texture2D = UTexture2D::CreateTransient(
						ImageWrapper->GetWidth(), ImageWrapper->GetHeight(), PF_R8G8B8A8);
					void* TextureData = Texture2D->PlatformData->Mips[0].BulkData.Lock(LOCK_READ_WRITE);
					FMemory::Memcpy(TextureData, UncompressedRGBA.GetData(), UncompressedRGBA.Num());
					// TextureData指针释放已经通过Unlock完成了。
					Texture2D->PlatformData->Mips[0].BulkData.Unlock();
					Texture2D->UpdateResource();
					return Texture2D.Get();
				}
			}

	return nullptr;
}

UTexture2D* UScreenShotSubsystem::LoadTexture2DFromLocal(FString& ImagePath, const bool bAbsPath) const
{
	TArray<FString> ImageFiles;
	FPaths::NormalizeDirectoryName(ImagePath);
	FString FinalPath = FString();
	if(bAbsPath)
	{
		FinalPath = ImagePath;
	}else
	{
		FinalPath = FPaths::ProjectContentDir() + ImagePath;
	}
	IFileManager& FileManager = IFileManager::Get();
	FileManager.FindFiles(ImageFiles, *FinalPath, true, true);
	if (ImageFiles.Num() > 0)
	{
		FString ImgPath = ImageFiles[0];
		TArray<uint8> CompressedData;
		if (FFileHelper::LoadFileToArray(CompressedData, *FinalPath))
		{
			EImageFormat ImageFormat = EImageFormat::Invalid;
			if (ImgPath.EndsWith(".png"))
			{
				ImageFormat = EImageFormat::PNG;
			}
			else if (ImgPath.EndsWith(".jpg") || ImgPath.EndsWith(".jpeg") || ImgPath.EndsWith(".PNG") || ImgPath.
				EndsWith(".JPEG"))
			{
				ImageFormat = EImageFormat::JPEG;
			}
			IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(
				FName("ImageWrapper"));
			TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(ImageFormat);
			if (ImageWrapper->SetCompressed(CompressedData.GetData(), CompressedData.Num()))
			{
				TArray<uint8> UncompressedRGBA;
				if (ImageWrapper->GetRaw(ERGBFormat::RGBA, 8, UncompressedRGBA))
				{
					TObjectPtr<UTexture2D> Texture2D = UTexture2D::CreateTransient(
						ImageWrapper->GetWidth(), ImageWrapper->GetHeight(), PF_R8G8B8A8);
					void* TextureData = Texture2D->PlatformData->Mips[0].BulkData.Lock(LOCK_READ_WRITE);
					FMemory::Memcpy(TextureData, UncompressedRGBA.GetData(), UncompressedRGBA.Num());
					// TextureData指针释放已经通过Unlock完成了。
					Texture2D->PlatformData->Mips[0].BulkData.Unlock();
					Texture2D->UpdateResource();
					return Texture2D.Get();
				}
			}
		}
	}
	return nullptr;
}

FString UScreenShotSubsystem::CaptureImage(FScreenShotParams& SSParams)
{
	if (!GEngine || !GEngine->GameViewport)
	{
		return FString();
	}
	const FString SavedFilePath = !SSParams.SavedFilePath.IsEmpty()
		                              ? SSParams.SavedFilePath
		                              : FString::Printf(TEXT("%s/Screenshots/"), *FPaths::ProjectContentDir());
	FString ReturnFilePath = FString::Printf(TEXT("%s/%s"), *SavedFilePath, *SSParams.Filename);
	if (SSParams.SSSuffix == EScreenShotSuffix::ESSS_JPG)
	{
		ReturnFilePath = FString::Printf(TEXT("%s%s"), *ReturnFilePath, *FString(".jpg"));
	}
	else if (SSParams.SSSuffix == EScreenShotSuffix::ESSS_PNG)
	{
		ReturnFilePath = FString::Printf(TEXT("%s%s"), *ReturnFilePath, *FString(".png"));
	}
	SSParams.ReturnFilePath = ReturnFilePath;
	DelegateHandle = UGameViewportClient::OnScreenshotCaptured().AddStatic(&OnScreenshotCompleteImage, SSParams);

	if (SSParams.ScreenShotType == EScreenShotType::ESST_Normal)
	{
		if (GEngine->GameViewport->Exec(nullptr, TEXT("Shot"), *GLog))
		{
			return ReturnFilePath;
		}
	}
	else if (SSParams.ScreenShotType == EScreenShotType::ESST_NormalWithUI)
	{
		if (GEngine->GameViewport->Exec(nullptr, TEXT("Shot showui"), *GLog))
		{
			return ReturnFilePath;
		}
	}
	else if (SSParams.ScreenShotType == EScreenShotType::ESST_HighResShot)
	{
		const FString Command = FString::Printf(
			TEXT("HighResShot %dx%d"), SSParams.Resolution.Width, SSParams.Resolution.Height);
		if (GEngine->GameViewport->Exec(nullptr, *Command, *GLog))
		{
			return ReturnFilePath;
		}
	}
	return FString();
}

void UScreenShotSubsystem::OnScreenshotCompleteImage(int32 InWidth, int32 InHeight, const TArray<FColor>& InColors,
                                                     const FScreenShotParams SSParams)
{
	IImageWrapperModule& ImageWrapperModule = FModuleManager::Get().LoadModuleChecked<IImageWrapperModule>(
		FName("ImageWrapper"));
	TSharedPtr<IImageWrapper> ImageWrapper = nullptr;
	if (SSParams.SSSuffix == EScreenShotSuffix::ESSS_JPG)
	{
		ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::JPEG);
	}
	else if (SSParams.SSSuffix == EScreenShotSuffix::ESSS_PNG)
	{
		ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);
	}
	if (ImageWrapper->SetRaw(InColors.GetData(), InColors.Num() * sizeof(FColor),
	                         InWidth, InHeight, ERGBFormat::BGRA, 8))
	{
		FFileHelper::SaveArrayToFile(ImageWrapper->GetCompressed(), *SSParams.ReturnFilePath);
	}
	UGameViewportClient::OnScreenshotCaptured().Remove(DelegateHandle);
}

TArray<FString> UScreenShotSubsystem::GetLocalAllFiles(FString& FilePath, const FString FileFix,
                                                       const EFixType FixType)
{
	TArray<FString> Files = TArray<FString>();
	FPaths::NormalizeDirectoryName(FilePath);
	IFileManager& FileManager = IFileManager::Get();
	const FString FinalPath = FilePath / TEXT("*");
	FileManager.FindFiles(Files, *FinalPath, true, true);
	TArray<FString> ReturnFiles = TArray<FString>();
	if (FixType == EFixType::EFT_Prefix)
	{
		for (const FString File : Files)
		{
			if(File.StartsWith(*FileFix))
			{
				ReturnFiles.Add(File);
			}
		}
	}
	else if (FixType == EFixType::EFT_Suffix)
	{
		for (const FString File : Files)
		{
			if(File.EndsWith(*FileFix))
			{
				ReturnFiles.Add(File);
			}
		}
	}
	return ReturnFiles;
}

UTexture2D* UScreenShotSubsystem::CreateQRCodeToTexture2D(const FString& URL, int32 PixelSize, FColor BackColor, FColor ForeColor) const
{
	// 1. 转换字符串为 UTF-8
	FTCHARToUTF8 Converter(*URL);
	const char* Utf8Text = Converter.Get();

	// 2. 生成二维码
	QRcode* QR = QRcode_encodeString(Utf8Text, 4, QR_ECLEVEL_H, QR_MODE_8, 1);
	if (!QR) return nullptr;

	const int32 QRWidth = QR->width;
	const int32 ImageSize = QRWidth * PixelSize;

	// 3. 创建颜色数组
	TArray<FColor> Pixels;
	Pixels.SetNumZeroed(ImageSize * ImageSize);

	for (int32 y = 0; y < QRWidth; ++y)
	{
		for (int32 x = 0; x < QRWidth; ++x)
		{
			const uint8 Module = QR->data[y * QRWidth + x] & 0x01;
			const FColor Color = Module ? ForeColor : BackColor;

			for (int32 dy = 0; dy < PixelSize; ++dy)
			{
				for (int32 dx = 0; dx < PixelSize; ++dx)
				{
					const int32 Px = x * PixelSize + dx;
					const int32 Py = y * PixelSize + dy;
					Pixels[Py * ImageSize + Px] = Color;
				}
			}
		}
	}

	// 4. 创建 UTexture2D
	UTexture2D* Texture = UTexture2D::CreateTransient(ImageSize, ImageSize, PF_B8G8R8A8);
	if (!Texture)
	{
		QRcode_free(QR);
		return nullptr;
	}

	// 5. 填充像素数据
	void* TextureData = Texture->PlatformData->Mips[0].BulkData.Lock(LOCK_READ_WRITE);
	FMemory::Memcpy(TextureData, Pixels.GetData(), Pixels.Num() * sizeof(FColor));
	Texture->PlatformData->Mips[0].BulkData.Unlock();

	// 6. 更新贴图资源
	Texture->UpdateResource();
	// Texture->MipGenSettings = TMGS_NoMipmaps;
	Texture->SRGB = true;

	// 7. 释放二维码资源
	QRcode_free(QR);
	return Texture;
}

void UScreenShotSubsystem::DeleteFolder(const FString FilePath)
{
	FPlatformFileManager::Get().Get().GetPlatformFile().DeleteDirectoryRecursively(*FilePath);
}

FString UScreenShotSubsystem::Base64URLEncode(const FString& InString)
{
	std::string StdString = TCHAR_TO_UTF8(*InString);
	static const char* base64_chars =
		"ABCDEFGHIJKLMNOPQRSTUVWXYZ"
		"abcdefghijklmnopqrstuvwxyz"
		"0123456789-_";

	std::string result;
	int val = 0, valb = -6;
	for (unsigned char c : StdString) {
		val = (val << 8) + c;
		valb += 8;
		while (valb >= 0) {
			result.push_back(base64_chars[(val >> valb) & 0x3F]);
			valb -= 6;
		}
	}
	if (valb > -6) {
		result.push_back(base64_chars[((val << 8) >> (valb + 8)) & 0x3F]);
	}
	return FString(UTF8_TO_TCHAR(result.c_str()));
}

FString UScreenShotSubsystem::Base64URLDecode(const FString& InString)
{
	std::string StdString = TCHAR_TO_UTF8(*InString);
	static const int B64index[256] = {
							-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, //   0-9
							-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, //  10-19
							-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, //  20-29
							-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, //  30-39
							-1,-1,-1,62,-1,-1,-1,63,52,53, //  40-49
							54,55,56,57,58,59,60,61,-1,-1, //  50-59
							-1,-1,-1,-1,-1, 0, 1, 2, 3, 4, //  60-69
							 5, 6, 7, 8, 9,10,11,12,13,14, //  70-79
							15,16,17,18,19,20,21,22,23,24, //  80-89
							25,-1,-1,-1,-1,-1,-1,26,27,28, //  90-99
							29,30,31,32,33,34,35,36,37,38, // 100-109
							39,40,41,42,43,44,45,46,47,48, // 110-119
							49,50,51,-1,-1,-1,-1,-1         // 120-127
	};

	std::vector<unsigned char> bytes;
	int val = 0, valb = -8;
	for (unsigned char c : StdString) {
		if (c >= 128 || B64index[c] == -1) break;
		val = (val <<	6) + B64index[c];
		valb += 6;
		if (valb >= 0) {
			bytes.push_back((val >> valb) & 0xFF);
			valb -= 8;
		}
	}

	return FString(UTF8_TO_TCHAR(std::string(bytes.begin(), bytes.end()).c_str()));
}
